#!/usr/bin/env python
# -*- coding: utf-8 -*-

from typing import List, Dict, Any, Optional, Tuple
import asyncio

from pymilvus import Collection
from pymilvus.exceptions import MilvusException

from app.core.config import settings
from app.core.exceptions import VectorDBError
from app.core.logging import get_logger, LoggerMixin
from app.core.utils import validate_collection_name
from app.repositories.milvus import MilvusClient

# 创建日志记录器
logger = get_logger("vector_repo")


class VectorRepository(LoggerMixin):
    """向量仓库
    
    管理Milvus中的向量数据
    """
    
    def __init__(self, milvus_client: MilvusClient):
        """初始化向量仓库
        
        Args:
            milvus_client: Milvus客户端
        """
        self.client = milvus_client
        self.dimension = settings.EMBEDDING_DIMENSION
    
    def _get_collection_name(self, project_id: str) -> str:
        """获取集合名称
        
        Args:
            project_id: 项目ID
            
        Returns:
            str: 集合名称
        """
        # 将项目ID转换为符合Milvus集合命名规范的格式
        safe_project_id = project_id.replace("-", "_")
        
        # 确保集合名称不以数字开头（Milvus集合名称要求）
        collection_prefix = f"p_{safe_project_id}" if safe_project_id[0].isdigit() else safe_project_id
        collection_name = f"train_{collection_prefix}"
        
        return collection_name
    
    async def ensure_collections(self, project_id: str) -> None:
        """确保项目的集合存在
        
        Args:
            project_id: 项目ID
            
        Raises:
            VectorDBError: 当创建集合失败时
        """
        # 创建单个训练数据集合
        collection_name = self._get_collection_name(project_id)
        await self.client.create_collection(
            collection_name=collection_name,
            dimension=self.dimension,
            description=f"Project {project_id} training data vectors"
        )
    
    async def delete_project_collections(self, project_id: str) -> None:
        """删除项目的集合
        
        Args:
            project_id: 项目ID
            
        Raises:
            VectorDBError: 当删除集合失败时
        """
        # 删除项目的训练数据集合
        collection_name = self._get_collection_name(project_id)
        await self.client.drop_collection(collection_name)
    
    async def insert_vector(
        self,
        project_id: str,
        data_type: str,
        content_id: int,
        vector: List[float],
        content_type: str = ""
    ) -> int:
        """插入向量
        
        Args:
            project_id: 项目ID
            data_type: 数据类型（现在作为content_type的一部分）
            content_id: 内容ID
            vector: 向量
            content_type: 内容类型，用于过滤
            
        Returns:
            int: 插入的向量ID
            
        Raises:
            VectorDBError: 当插入失败时
        """
        # 获取集合名称
        collection_name = self._get_collection_name(project_id)
        
        # 如果没有提供content_type，使用data_type作为content_type
        if not content_type:
            content_type = data_type
        
        try:
            # 获取集合
            collection = await self.client.get_collection(collection_name)
            
            # 准备数据
            data = [
                [content_id],     # content_id
                [content_type],   # data_type
                [vector]          # embedding
            ]
            
            # 插入数据
            loop = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                None,
                lambda: collection.insert(data)
            )
            
            # 返回插入的ID
            return result.primary_keys[0]
            
        except Exception as e:
            error_msg = f"插入向量失败: {str(e)}"
            self.logger.error(error_msg)
            raise VectorDBError(error_msg)
    
    async def delete_vector(
        self,
        project_id: str,
        data_type: str,
        content_id: int
    ) -> None:
        """删除向量
        
        Args:
            project_id: 项目ID
            data_type: 数据类型（现在仅用于日志记录）
            content_id: 内容ID
            
        Raises:
            VectorDBError: 当删除失败时
        """
        # 获取集合名称
        collection_name = self._get_collection_name(project_id)
        
        try:
            # 检查集合是否存在
            exists = await self.client.collection_exists(collection_name)
            if not exists:
                self.logger.warning(f"集合不存在，无需删除: {collection_name}")
                return
                
            # 获取集合
            collection = await self.client.get_collection(collection_name)
            
            # 构建表达式
            expr = f"content_id == {content_id}"
            
            # 删除数据
            loop = asyncio.get_event_loop()
            await loop.run_in_executor(
                None,
                lambda: collection.delete(expr)
            )
            
        except Exception as e:
            error_msg = f"删除向量失败: {str(e)}"
            self.logger.error(error_msg)
            raise VectorDBError(error_msg)
    
    async def search_similar(
        self,
        project_id: str,
        data_type: str,
        vector: List[float],
        top_k: int = 5,
        content_type: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """搜索相似向量
        
        Args:
            project_id: 项目ID
            data_type: 数据类型（现在用于过滤，如果没有提供content_type）
            vector: 查询向量
            top_k: 返回的最大结果数
            content_type: 内容类型，用于过滤
            
        Returns:
            List[Dict[str, Any]]: 相似向量列表
            
        Raises:
            VectorDBError: 当搜索失败时
        """
        # 获取集合名称
        collection_name = self._get_collection_name(project_id)
        
        # 如果没有提供content_type，使用data_type作为content_type
        if not content_type:
            content_type = data_type
        
        try:
            # 检查集合是否存在
            exists = await self.client.collection_exists(collection_name)
            if not exists:
                self.logger.warning(f"集合不存在: {collection_name}")
                return []
            
            # 获取集合
            collection = await self.client.get_collection(collection_name)
            
            # 搜索参数
            search_params = {
                "metric_type": "COSINE",
                "params": {"ef": 32}
            }
            
            # 构建表达式（如果提供了内容类型）
            expr = None
            if content_type:
                expr = f'data_type == "{content_type}"'
                
            # 执行搜索
            loop = asyncio.get_event_loop()
            results = await loop.run_in_executor(
                None,
                lambda: collection.search(
                    data=[vector],
                    anns_field="embedding",
                    param=search_params,
                    limit=top_k,
                    expr=expr,
                    output_fields=["content_id", "data_type"]
                )
            )
            
            # 处理结果
            similar_items = []
            for hits in results:
                for hit in hits:
                    similar_items.append({
                        "content_id": hit.entity.get("content_id"),
                        "data_type": hit.entity.get("data_type"),
                        "distance": hit.distance,
                        "score": 1.0 - hit.distance  # 转换为相似度分数
                    })
            
            return similar_items
            
        except Exception as e:
            error_msg = f"搜索相似向量失败: {str(e)}"
            self.logger.error(error_msg)
            raise VectorDBError(error_msg)